{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "(tir-transform)=\n",
        "# 元张量函数变换\n",
        "\n",
        "在本节中，将介绍编译流程的主要组成部分——元张量函数变换。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-cell"
        ]
      },
      "outputs": [],
      "source": [
        "import tvm\n",
        "from tvm.script import ir as I\n",
        "from tvm.script import tir as T\n",
        "\n",
        "\n",
        "@I.ir_module\n",
        "class MyModule:\n",
        "    @T.prim_func\n",
        "    def main(\n",
        "        A: T.Buffer((128, 128), \"float32\"),\n",
        "        B: T.Buffer((128, 128), \"float32\"),\n",
        "        C: T.Buffer((128, 128), \"float32\"),\n",
        "    ):\n",
        "        T.func_attr({\"tir.noalias\": T.bool(True)})\n",
        "        Y = T.alloc_buffer((128, 128))\n",
        "        for i, j, k in T.grid(128, 128, 128):\n",
        "            with T.block(\"Y\"):\n",
        "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
        "                with T.init():\n",
        "                    Y[vi, vj] = T.float32(0)\n",
        "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
        "        for i, j in T.grid(128, 128):\n",
        "            with T.block(\"C\"):\n",
        "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
        "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "在我们转换函数之前，让我们先评估原始实现的性能。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-cell"
        ]
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Execution time summary:\n",
            " mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  \n",
            "   2.1428       2.1428       2.1428       2.1428       0.0000                  \n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "\n",
        "a_np = np.random.uniform(size=(128, 128)).astype(\"float32\")\n",
        "b_np = np.random.uniform(size=(128, 128)).astype(\"float32\")\n",
        "c_np = a_np @ b_np\n",
        "\n",
        "a_nd = tvm.runtime.tensor(a_np)\n",
        "b_nd = tvm.runtime.tensor(b_np)\n",
        "c_nd = tvm.runtime.tensor(np.zeros((128, 128), dtype=\"float32\"))\n",
        "\n",
        "\n",
        "def evaluate(mod: tvm.IRModule):\n",
        "    lib = tvm.build(mod, target=\"llvm\")\n",
        "    # check correctness\n",
        "    lib(a_nd, b_nd, c_nd)\n",
        "    np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5)\n",
        "    # evaluate performance\n",
        "    f_timer = lib.time_evaluator(\"main\", tvm.cpu())\n",
        "    print(f_timer(a_nd, b_nd, c_nd))\n",
        "\n",
        "\n",
        "evaluate(MyModule)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 初始化调度\n",
        "\n",
        "使用提供的 **MyModule** 作为输入来建立 Schedule 辅助类，以此启动代码转换的过程。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sch = tvm.tir.Schedule(MyModule)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 循环平铺(Tiling)\n",
        "接下来，我们执行必要的操作来获取块 **Y** 及其关联循环的引用。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "block_Y = sch.get_block(\"Y\")\n",
        "i, j, k = sch.get_loops(block_Y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "我们现在开始执行转换。第一个修改是将循环``j``拆分为两个独立的循环，内层循环的长度为4。必须了解的是，转换过程是分步进行的；因此，如果无意中再次执行块，将会出现一个错误，指出变量``j``不存在。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "j0, j1 = sch.split(j, factors=[None, 8])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "转换的结果可以被检查，因为它被保留在``sch.mod``中。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        Y <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>))\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j_0, j_1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">128</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y&quot;</span>):\n",
              "                vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1)\n",
              "                vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">128</span>, k)\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[vi, vk], B[vk, vj])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "                Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                vi, vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i, j])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
              "                C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(Y[vi, vj], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sch.mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "在初始转换阶段之后，生成了两个补充循环``j_0``和``j_1``，它们的范围分别为32和4。接下来的动作是重新排序这两个循环。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        Y <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>))\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j_0, k, j_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">8</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y&quot;</span>):\n",
              "                vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1)\n",
              "                vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">128</span>, k)\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[vi, vk], B[vk, vj])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "                Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                vi, vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i, j])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
              "                C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(Y[vi, vj], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Execution time summary:\n",
            " mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  \n",
            "   1.1268       1.1268       1.1268       1.1268       0.0000                  \n"
          ]
        }
      ],
      "source": [
        "sch.reorder(j0, k, j1)\n",
        "sch.mod.show()\n",
        "evaluate(sch.mod)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 利用局部性\n",
        "接下来，我们将执行两个额外的转换步骤以获得不同的变体。首先，使用名为 **reverse_compute_at** 的原语将块 **C** 移动到 **Y** 的内层循环中。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        Y <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>))\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">16</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> k, j_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1)\n",
              "                    vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">128</span>, k)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[vi, vk], B[vk, vj])\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
              "                    <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                        Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ax0 <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> ax0)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj])\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
              "                    C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(Y[vi, vj], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "block_C = sch.get_block(\"C\")\n",
        "sch.reverse_compute_at(block_C, j0)\n",
        "sch.mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 重写约简\n",
        "到目前为止，约简初始化和更新步骤一直保留在同一个块体中。这种混合形式便于循环转换，因为初始化和更新的外部循环``i``, ``j``通常需要保持同步。\n",
        "\n",
        "在循环转换之后，我们可以使用**decompose_reduction**原语将Y元素的初始化与约简更新分开。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        Y <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>))\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">16</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> j_1_init <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y_init&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1_init)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
              "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> k, j_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y_update&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1)\n",
              "                    vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">128</span>, k)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj], A[vi, vk], B[vk, vj])\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
              "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ax0 <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> ax0)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj])\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
              "                    C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(Y[vi, vj], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Execution time summary:\n",
            " mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  \n",
            "   0.4294       0.4294       0.4294       0.4294       0.0000                  \n"
          ]
        }
      ],
      "source": [
        "sch.decompose_reduction(block_Y, k)\n",
        "sch.mod.show()\n",
        "evaluate(sch.mod)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 跟踪转换\n",
        "TensorIR调度是一种过程语言，转换是以逐步方式执行的。我们可以通过打印调度或调度的历史来跟踪转换。\n",
        "\n",
        "我们已经通过打印``sch.mod``看到了调度。我们还可以通过``sch.trace``打印调度的历史。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm import tir</span>\n",
              "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">apply_trace</span>(sch: tir<span style=\"color: #AA22FF; font-weight: bold\">.</span>Schedule) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "  b0 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;Y&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
              "  l1, l2, l3 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_loops(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0)\n",
              "  l4, l5 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>split(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l2, factors<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #008000\">8</span>], preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, disable_predication<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>)\n",
              "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>reorder(l4, l3, l5)\n",
              "  b6 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;C&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
              "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>reverse_compute_at(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b6, loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l4, preserve_unit_loops<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>, index<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">1</span>)\n",
              "  b7 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>decompose_reduction(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0, loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l3)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sch.trace.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "或者，我们可以结合历史跟踪输出IRModule。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "collapsed": false,
        "tags": [
          "hide-output"
        ]
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        Y <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">128</span>))\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">16</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> j_1_init <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y_init&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1_init)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
              "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> k, j_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;Y_update&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_1)\n",
              "                    vk <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">128</span>, k)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj], A[vi, vk], B[vk, vj])\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(Y[vi, vj])\n",
              "                    Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> Y[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[vi, vk] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[vk, vj]\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ax0 <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">8</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                    vj <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, j_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> ax0)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(Y[vi, vj])\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi, vj])\n",
              "                    C[vi, vj] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(Y[vi, vj], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>))\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm import tir</span>\n",
              "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">apply_trace</span>(sch: tir<span style=\"color: #AA22FF; font-weight: bold\">.</span>Schedule) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "  b0 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;Y&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
              "  l1, l2, l3 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_loops(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0)\n",
              "  l4, l5 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>split(loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l2, factors<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #008000\">8</span>], preserve_unit_iters<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, disable_predication<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>)\n",
              "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>reorder(l4, l3, l5)\n",
              "  b6 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>get_block(name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;C&quot;</span>, func_name<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;main&quot;</span>)\n",
              "  sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>reverse_compute_at(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b6, loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l4, preserve_unit_loops<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>, index<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">1</span>)\n",
              "  b7 <span style=\"color: #AA22FF; font-weight: bold\">=</span> sch<span style=\"color: #AA22FF; font-weight: bold\">.</span>decompose_reduction(block<span style=\"color: #AA22FF; font-weight: bold\">=</span>b0, loop<span style=\"color: #AA22FF; font-weight: bold\">=</span>l3)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sch.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "ai",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
